{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 5: Equation solving"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Robert Johansson\n",
    "\n",
    "Source code listings for [Numerical Python - Scientific Computing and Data Science Applications with Numpy, SciPy and Matplotlib](https://www.apress.com/us/book/9781484242452) (ISBN 978-1-484242-45-2)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "matplotlib.rcParams['mathtext.fontset'] = 'stix'\n",
    "matplotlib.rcParams['font.family'] = 'serif'\n",
    "matplotlib.rcParams['font.sans-serif'] = 'stix'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from scipy import linalg as la"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from scipy import optimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sympy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "sympy.init_printing()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "#import matplotlib as mpl\n",
    "#mpl.rcParams[\"font.family\"] = \"serif\"\n",
    "#mpl.rcParams[\"font.size\"] = \"12\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from __future__ import division"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear Algebra - Linear Equation Systems"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "2 x_1 + 3 x_2 = 4\n",
    "$$\n",
    "\n",
    "$$\n",
    "5 x_1 + 4 x_2 = 3\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "x1 = np.linspace(-4, 2, 100)\n",
    "\n",
    "x2_1 = (4 - 2 * x1)/3\n",
    "x2_2 = (3 - 5 * x1)/4\n",
    "\n",
    "ax.plot(x1, x2_1, 'r', lw=2, label=r\"$2x_1+3x_2-4=0$\")\n",
    "ax.plot(x1, x2_2, 'b', lw=2, label=r\"$5x_1+4x_2-3=0$\")\n",
    "\n",
    "A = np.array([[2, 3], [5, 4]])\n",
    "b = np.array([4, 3])\n",
    "x = la.solve(A, b)\n",
    "\n",
    "ax.plot(x[0], x[1], 'ko', lw=2)\n",
    "ax.annotate(\"The intersection point of\\nthe two lines is the solution\\nto the equation system\",\n",
    "            xy=(x[0], x[1]), xycoords='data',\n",
    "            xytext=(-120, -75), textcoords='offset points', \n",
    "            arrowprops=dict(arrowstyle=\"->\", connectionstyle=\"arc3, rad=-.3\"))\n",
    "\n",
    "ax.set_xlabel(r\"$x_1$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$x_2$\", fontsize=18)\n",
    "ax.legend();\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch5-linear-systems-simple.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Symbolic approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "A = sympy.Matrix([[2, 3], [5, 4]])\n",
    "b = sympy.Matrix([4, 3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "A.rank()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "A.condition_number()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "sympy.N(_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "A.norm()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "L, U, P = A.LUdecomposition()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "U"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "L * U"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "x = A.solve(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Numerical approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "A = np.array([[2, 3], [5, 4]])\n",
    "b = np.array([4, 3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "np.linalg.matrix_rank(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "np.linalg.cond(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "np.linalg.norm(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "P, L, U = la.lu(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "U"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "np.dot(L, U)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "np.dot(P, np.dot(L, U))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "P.dot(L.dot(U))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "la.solve(A, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example : rank and condition numbers -> numerical errors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "p = sympy.symbols(\"p\", positive=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "A = sympy.Matrix([[1, sympy.sqrt(p)], [1, 1/sympy.sqrt(p)]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "b = sympy.Matrix([1, 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "sympy.simplify(A.solve(b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Symbolic problem specification\n",
    "p = sympy.symbols(\"p\", positive=True)\n",
    "A = sympy.Matrix([[1, sympy.sqrt(p)], [1, 1/sympy.sqrt(p)]])\n",
    "b = sympy.Matrix([1, 2])\n",
    "\n",
    "# Solve symbolically\n",
    "x_sym_sol = A.solve(b)\n",
    "x_sym_sol.simplify()\n",
    "x_sym_sol\n",
    "Acond = A.condition_number().simplify()\n",
    "\n",
    "# Function for solving numerically\n",
    "AA = lambda p: np.array([[1, np.sqrt(p)], [1, 1/np.sqrt(p)]])\n",
    "bb = np.array([1, 2])\n",
    "x_num_sol = lambda p: np.linalg.solve(AA(p), bb)\n",
    "\n",
    "# Graph the difference between the symbolic (exact) and numerical results.\n",
    "p_vec = np.linspace(0.9, 1.1, 200)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "for n in range(2):\n",
    "    x_sym = np.array([x_sym_sol[n].subs(p, pp).evalf() for pp in p_vec])\n",
    "    x_num = np.array([x_num_sol(pp)[n] for pp in p_vec])\n",
    "    axes[0].plot(p_vec, (x_num - x_sym)/x_sym, 'k')\n",
    "axes[0].set_title(\"Error in solution\\n(numerical - symbolic)/symbolic\")\n",
    "axes[0].set_xlabel(r'$p$', fontsize=18)\n",
    "\n",
    "axes[1].plot(p_vec, [Acond.subs(p, pp).evalf() for pp in p_vec])\n",
    "axes[1].set_title(\"Condition number\")\n",
    "axes[1].set_xlabel(r'$p$', fontsize=18)\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch5-linear-systems-condition-number.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rectangular systems"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Underdetermined"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "unknown = sympy.symbols(\"x, y, z\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "A = sympy.Matrix([[1, 2, 3], [4, 5, 6]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "x = sympy.Matrix(unknown)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "b = sympy.Matrix([7, 8])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "AA = A * x - b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "sympy.solve(A*x - b, unknown)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Overdetermined: least squares"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "np.random.seed(1234)\n",
    "\n",
    "# define true model parameters\n",
    "x = np.linspace(-1, 1, 100)\n",
    "a, b, c = 1, 2, 3\n",
    "y_exact = a + b * x + c * x**2\n",
    "\n",
    "# simulate noisy data points\n",
    "m = 100\n",
    "X = 1 - 2 * np.random.rand(m)\n",
    "Y = a + b * X + c * X**2 + np.random.randn(m)\n",
    "\n",
    "# fit the data to the model using linear least square\n",
    "A = np.vstack([X**0, X**1, X**2])  # see np.vander for alternative\n",
    "sol, r, rank, sv = la.lstsq(A.T, Y)\n",
    "y_fit = sol[0] + sol[1] * x + sol[2] * x**2\n",
    "fig, ax = plt.subplots(figsize=(12, 4))\n",
    "\n",
    "ax.plot(X, Y, 'go', alpha=0.5, label='Simulated data')\n",
    "ax.plot(x, y_exact, 'k', lw=2, label='True value $y = 1 + 2x + 3x^2$')\n",
    "ax.plot(x, y_fit, 'b', lw=2, label='Least square fit')\n",
    "ax.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$y$\", fontsize=18)\n",
    "ax.legend(loc=2);\n",
    "\n",
    "fig.savefig('ch5-linear-systems-least-square.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# fit the data to the model using linear least square: \n",
    "# 1st order polynomial\n",
    "A = np.vstack([X**n for n in range(2)])\n",
    "sol, r, rank, sv = la.lstsq(A.T, Y)\n",
    "y_fit1 = sum([s * x**n for n, s in enumerate(sol)])\n",
    "\n",
    "# 15th order polynomial\n",
    "A = np.vstack([X**n for n in range(16)])\n",
    "sol, r, rank, sv = la.lstsq(A.T, Y)\n",
    "y_fit15 = sum([s * x**n for n, s in enumerate(sol)])\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(12, 4))\n",
    "ax.plot(X, Y, 'go', alpha=0.5, label='Simulated data')\n",
    "ax.plot(x, y_exact, 'k', lw=2, label='True value $y = 1 + 2x + 3x^2$')\n",
    "ax.plot(x, y_fit1, 'b', lw=2, label='Least square fit [1st order]')\n",
    "ax.plot(x, y_fit15, 'm', lw=2, label='Least square fit [15th order]')\n",
    "ax.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$y$\", fontsize=18)\n",
    "ax.legend(loc=2);\n",
    "\n",
    "fig.savefig('ch5-linear-systems-least-square-2.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Eigenvalue problems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "eps, delta = sympy.symbols(\"epsilon, delta\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "H = sympy.Matrix([[eps, delta], [delta, -eps]])\n",
    "H"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "eval1, eval2 = H.eigenvals()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "eval1, eval2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "H.eigenvects()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "(eval1, _, evec1), (eval2, _, evec2) = H.eigenvects()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "sympy.simplify(evec1[0].T * evec2[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "A = np.array([[1, 3, 5], [3, 5, 3], [5, 3, 9]])\n",
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "evals, evecs = la.eig(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "evals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "evecs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "la.eigvalsh(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nonlinear equations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Univariate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "x = np.linspace(-2, 2, 1000)\n",
    "\n",
    "# four examples of nonlinear functions\n",
    "f1 = x**2 - x - 1\n",
    "f2 = x**3 - 3 * np.sin(x)\n",
    "f3 = np.exp(x) - 2\n",
    "f4 = 1 - x**2 + np.sin(50 / (1 + x**2))\n",
    "\n",
    "# plot each function\n",
    "fig, axes = plt.subplots(1, 4, figsize=(12, 3), sharey=True)\n",
    "\n",
    "for n, f in enumerate([f1, f2, f3, f4]):\n",
    "    axes[n].plot(x, f, lw=1.5)\n",
    "    axes[n].axhline(0, ls=':', color='k')\n",
    "    axes[n].set_ylim(-5, 5)\n",
    "    axes[n].set_xticks([-2, -1, 0, 1, 2])\n",
    "    axes[n].set_xlabel(r'$x$', fontsize=18)\n",
    "\n",
    "axes[0].set_ylabel(r'$f(x)$', fontsize=18)\n",
    "\n",
    "titles = [r'$f(x)=x^2-x-1$', r'$f(x)=x^3-3\\sin(x)$',\n",
    "          r'$f(x)=\\exp(x)-2$', r'$f(x)=\\sin\\left(50/(1+x^2)\\right)+1-x^2$']\n",
    "for n, title in enumerate(titles):\n",
    "    axes[n].set_title(title)\n",
    "    \n",
    "fig.tight_layout()\n",
    "fig.savefig('ch5-nonlinear-plot-equations.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Symbolic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "x, a, b, c = sympy.symbols(\"x, a, b, c\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "e = a + b*x + c*x**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "sol1, sol2 = sympy.solve(e, x)\n",
    "\n",
    "sol1, sol2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "e.subs(x, sol1).expand()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "e.subs(x, sol2).expand()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "e = a * sympy.cos(x) - b * sympy.sin(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sol1, sol2 = sympy.solve(e, x)\n",
    "\n",
    "sol1, sol2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "e.subs(x, sympy.atan(a/b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "e.subs(x, sol1).simplify()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "e.subs(x, sol2).simplify()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "sympy.solve(sympy.sin(x)-x, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bisection method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# define a function, desired tolerance and starting interval [a, b]\n",
    "f = lambda x: np.exp(x) - 2\n",
    "tol = 0.1\n",
    "a, b = -2, 2\n",
    "x = np.linspace(-2.1, 2.1, 1000)\n",
    "\n",
    "# graph the function f\n",
    "fig, ax = plt.subplots(1, 1, figsize=(12, 4))\n",
    "\n",
    "ax.plot(x, f(x), lw=1.5)\n",
    "ax.axhline(0, ls=':', color='k')\n",
    "ax.set_xticks([-2, -1, 0, 1, 2])\n",
    "ax.set_xlabel(r'$x$', fontsize=18)\n",
    "ax.set_ylabel(r'$f(x)$', fontsize=18)\n",
    "\n",
    "# find the root using the bisection method and visualize\n",
    "# the steps in the method in the graph\n",
    "fa, fb = f(a), f(b)\n",
    "\n",
    "ax.plot(a, fa, 'ko')\n",
    "ax.plot(b, fb, 'ko')\n",
    "ax.text(a, fa + 0.5, r\"$a$\", ha='center', fontsize=18)\n",
    "ax.text(b, fb + 0.5, r\"$b$\", ha='center', fontsize=18)\n",
    "\n",
    "n = 1\n",
    "while b - a > tol:\n",
    "    m = a + (b - a)/2\n",
    "    fm = f(m)\n",
    "\n",
    "    ax.plot(m, fm, 'ko')\n",
    "    ax.text(m, fm - 0.5, r\"$m_%d$\" % n, ha='center')\n",
    "    n += 1\n",
    "    \n",
    "    if np.sign(fa) == np.sign(fm):\n",
    "        a, fa = m, fm\n",
    "    else:\n",
    "        b, fb = m, fm\n",
    "\n",
    "ax.plot(m, fm, 'r*', markersize=10)\n",
    "ax.annotate(\"Root approximately at %.3f\" % m,\n",
    "            fontsize=14, family=\"serif\",\n",
    "            xy=(a, fm), xycoords='data',\n",
    "            xytext=(-150, +50), textcoords='offset points', \n",
    "            arrowprops=dict(arrowstyle=\"->\", connectionstyle=\"arc3, rad=-.5\"))\n",
    "\n",
    "ax.set_title(\"Bisection method\")\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch5-nonlinear-bisection.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# define a function, desired tolerance and starting point xk\n",
    "tol = 0.01\n",
    "xk = 2\n",
    "\n",
    "s_x = sympy.symbols(\"x\")\n",
    "s_f = sympy.exp(s_x) - 2\n",
    "\n",
    "f = lambda x: sympy.lambdify(s_x, s_f, 'numpy')(x)\n",
    "fp = lambda x: sympy.lambdify(s_x, sympy.diff(s_f, s_x), 'numpy')(x)\n",
    "\n",
    "x = np.linspace(-1, 2.1, 1000)\n",
    "\n",
    "# setup a graph for visualizing the root finding steps\n",
    "fig, ax = plt.subplots(1, 1, figsize=(12,4))\n",
    "\n",
    "ax.plot(x, f(x))\n",
    "ax.axhline(0, ls=':', color='k')\n",
    "\n",
    "# repeat Newton's method until convergence to the desired tolerance has been reached\n",
    "n = 0\n",
    "while f(xk) > tol:\n",
    "    xk_new = xk - f(xk) / fp(xk)\n",
    "\n",
    "    ax.plot([xk, xk], [0, f(xk)], color='k', ls=':')\n",
    "    ax.plot(xk, f(xk), 'ko')\n",
    "    ax.text(xk, -.5, r'$x_%d$' % n, ha='center')\n",
    "    ax.plot([xk, xk_new], [f(xk), 0], 'k-')\n",
    "\n",
    "    xk = xk_new\n",
    "    n += 1\n",
    "\n",
    "ax.plot(xk, f(xk), 'r*', markersize=15)\n",
    "ax.annotate(\"Root approximately at %.3f\" % xk,\n",
    "            fontsize=14, family=\"serif\",\n",
    "            xy=(xk, f(xk)), xycoords='data',\n",
    "            xytext=(-150, +50), textcoords='offset points', \n",
    "            arrowprops=dict(arrowstyle=\"->\", connectionstyle=\"arc3, rad=-.5\"))\n",
    "\n",
    "ax.set_title(\"Newton's method\")\n",
    "ax.set_xticks([-1, 0, 1, 2])\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch5-nonlinear-newton.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `scipy.optimize` functions for root-finding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "optimize.bisect(lambda x: np.exp(x) - 2, -2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "optimize.newton(lambda x: np.exp(x) - 2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "x_root_guess = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "f = lambda x: np.exp(x) - 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "fprime = lambda x: np.exp(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "optimize.newton(f, x_root_guess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "optimize.newton(f, x_root_guess, fprime=fprime)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "optimize.brentq(lambda x: np.exp(x) - 2, -2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "optimize.brenth(lambda x: np.exp(x) - 2, -2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "optimize.ridder(lambda x: np.exp(x) - 2, -2, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multivariate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return [x[1] - x[0]**3 - 2 * x[0]**2 + 1, x[1] + x[0]**2 - 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "optimize.fsolve(f, [1, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def f_jacobian(x):\n",
    "    return [[-3*x[0]**2-4*x[0], 1], [2*x[0], 1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "optimize.fsolve(f, [1, 1], fprime=f_jacobian)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "x, y = sympy.symbols(\"x, y\")\n",
    "\n",
    "f_mat = sympy.Matrix([y - x**3 -2*x**2 + 1, y + x**2 - 1])\n",
    "f_mat.jacobian(sympy.Matrix([x, y]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "#def f(x):\n",
    "#    return [x[1] - x[0]**3 - 2 * x[0]**2 + 1, x[1] + x[0]**2 - 1]\n",
    "\n",
    "x = np.linspace(-3, 2, 5000)\n",
    "y1 = x**3 + 2 * x**2 -1\n",
    "y2 = -x**2 + 1\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "ax.plot(x, y1, 'b', lw=1.5, label=r'$y = x^3 + 2x^2 - 1$')\n",
    "ax.plot(x, y2, 'g', lw=1.5, label=r'$y = -x^2 + 1$')\n",
    "\n",
    "x_guesses = [[-2, 2], [1, -1], [-2, -5]]\n",
    "for x_guess in x_guesses:\n",
    "    sol = optimize.fsolve(f, x_guess)\n",
    "    ax.plot(sol[0], sol[1], 'r*', markersize=15)\n",
    "\n",
    "    ax.plot(x_guess[0], x_guess[1], 'ko')\n",
    "    ax.annotate(\"\", xy=(sol[0], sol[1]), xytext=(x_guess[0], x_guess[1]),\n",
    "                arrowprops=dict(arrowstyle=\"->\", linewidth=2.5))\n",
    "    \n",
    "ax.legend(loc=0)\n",
    "ax.set_xlabel(r'$x$', fontsize=18)\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch5-nonlinear-system.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "optimize.broyden2(f, x_guesses[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return [x[1] - x[0]**3 - 2 * x[0]**2 + 1,\n",
    "            x[1] + x[0]**2 - 1]\n",
    "\n",
    "x = np.linspace(-3, 2, 5000)\n",
    "y1 = x**3 + 2 * x**2 -1\n",
    "y2 = -x**2 + 1\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "ax.plot(x, y1, 'k', lw=1.5, label=r'$y = x^3 + 2x^2 - 1$')\n",
    "ax.plot(x, y2, 'k', lw=1.5, label=r'$y = -x^2 + 1$')\n",
    "\n",
    "sol1 = optimize.fsolve(f, [-2,  2])\n",
    "sol2 = optimize.fsolve(f, [ 1, -1])\n",
    "sol3 = optimize.fsolve(f, [-2, -5])\n",
    "sols = [sol1, sol2, sol3]\n",
    "colors = ['r', 'b', 'g']\n",
    "for idx, s in enumerate(sols):\n",
    "    ax.plot(s[0], s[1], colors[idx]+'*', markersize=15)\n",
    "\n",
    "for m in np.linspace(-4, 3, 80):\n",
    "    for n in np.linspace(-15, 15, 40):\n",
    "        x_guess = [m, n]\n",
    "        sol = optimize.fsolve(f, x_guess)\n",
    "        idx = (abs(sols - sol)**2).sum(axis=1).argmin()\n",
    "        ax.plot(x_guess[0], x_guess[1], colors[idx]+'.')\n",
    "    \n",
    "ax.set_xlabel(r'$x$', fontsize=18)\n",
    "ax.set_xlim(-4, 3)\n",
    "ax.set_ylim(-15, 15)\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch5-nonlinear-system-map.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return [x[1] - x[0]**3 - 2 * x[0]**2 + 1,\n",
    "            x[1] + x[0]**2 - 1]\n",
    "\n",
    "x = np.linspace(-3, 2, 5000)\n",
    "y1 = x**3 + 2 * x**2 -1\n",
    "y2 = -x**2 + 1\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "ax.plot(x, y1, 'k', lw=1.5, label=r'$y = x^3 + 2x^2 - 1$')\n",
    "ax.plot(x, y2, 'k', lw=1.5, label=r'$y = -x^2 + 1$')\n",
    "\n",
    "sol1 = optimize.fsolve(f, [-2,  2])\n",
    "sol2 = optimize.fsolve(f, [ 1, -1])\n",
    "sol3 = optimize.fsolve(f, [-2, -5])\n",
    "\n",
    "for idx, s in enumerate([sol1, sol2, sol3]):\n",
    "    ax.plot(s[0], s[1], colors[idx]+'*', markersize=15)\n",
    "\n",
    "colors = ['r', 'b', 'g']\n",
    "for m in np.linspace(-4, 3, 80):\n",
    "    for n in np.linspace(-15, 15, 40):\n",
    "        x_guess = [m, n]\n",
    "        sol = optimize.fsolve(f, x_guess)\n",
    "\n",
    "        for idx, s in enumerate([sol1, sol2, sol3]):\n",
    "            if abs(s-sol).max() < 1e-8:\n",
    "                # ax.plot(sol[0], sol[1], colors[idx]+'*', markersize=15)\n",
    "                ax.plot(x_guess[0], x_guess[1], colors[idx]+'.')\n",
    "    \n",
    "ax.set_xlabel(r'$x$', fontsize=18)\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch5-nonlinear-system-map.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext version_information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%version_information sympy, scipy, numpy, matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
